{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7ac69fdc",
   "metadata": {},
   "source": [
    "# Tutorial 06, case 3a: advection diffusion reaction control problem with Neumann control\n",
    "\n",
    "In this tutorial we solve the optimal control problem\n",
    "\n",
    "$$\\min J(y, u) = \\frac{1}{2} \\int_{\\Omega} (y - y_d)^2 dx + \\frac{\\alpha}{2} \\int_{\\Gamma_2} u^2 ds$$\n",
    "s.t.\n",
    "$$\\begin{cases}\n",
    "- \\epsilon \\Delta y + \\beta \\cdot \\nabla y + \\sigma y = f      & \\text{in } \\Omega\\\\\n",
    "                                \\epsilon \\partial_n y = 0      & \\text{on } \\Gamma_1\\\\\n",
    "                                \\epsilon \\partial_n y = u      & \\text{on } \\Gamma_2\\\\\n",
    "                                \\epsilon \\partial_n y = 0      & \\text{on } \\Gamma_3\\\\\n",
    "                                                    y = 0      & \\text{on } \\Gamma_4\n",
    "\\end{cases}$$\n",
    "\n",
    "where\n",
    "$$\\begin{align*}\n",
    "& \\Omega               & \\text{unit square}\\\\\n",
    "& \\Gamma_1             & \\text{bottom boundary of the square}\\\\\n",
    "& \\Gamma_2             & \\text{left boundary of the square}\\\\\n",
    "& \\Gamma_3             & \\text{top boundary of the square}\\\\\n",
    "& \\Gamma_4             & \\text{right boundary of the square}\\\\\n",
    "& u \\in L^2(\\Gamma_2)  & \\text{control variable}\\\\\n",
    "& y \\in H^1(\\Omega)    & \\text{state variable}\\\\\n",
    "& \\alpha > 0           & \\text{penalization parameter}\\\\\n",
    "& y_d                  & \\text{desired state}\\\\\n",
    "& f                    & \\text{forcing term}\n",
    "\\end{align*}$$\n",
    "using an adjoint formulation solved by a one shot approach"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae42cc87",
   "metadata": {},
   "outputs": [],
   "source": [
    "import dolfinx.fem\n",
    "import dolfinx.fem.petsc\n",
    "import dolfinx.io\n",
    "import dolfinx.mesh\n",
    "import mpi4py.MPI\n",
    "import numpy as np\n",
    "import numpy.typing\n",
    "import petsc4py.PETSc\n",
    "import ufl\n",
    "import viskex"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "408b709d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import multiphenicsx.fem\n",
    "import multiphenicsx.fem.petsc"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a6460550",
   "metadata": {},
   "source": [
    "### Mesh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5915a74e",
   "metadata": {},
   "outputs": [],
   "source": [
    "mesh = dolfinx.mesh.create_unit_square(mpi4py.MPI.COMM_WORLD, 32, 32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b60b7b2-f7b0-48eb-b807-2050d057124b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create connectivities required by the rest of the code\n",
    "mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "787330e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def bottom(x: np.typing.NDArray[np.float64]) -> np.typing.NDArray[np.bool_]:\n",
    "    \"\"\"Condition that defines the bottom boundary.\"\"\"\n",
    "    return abs(x[1] - 0.) < np.finfo(float).eps  # type: ignore[no-any-return]\n",
    "\n",
    "\n",
    "def left(x: np.typing.NDArray[np.float64]) -> np.typing.NDArray[np.bool_]:\n",
    "    \"\"\"Condition that defines the left boundary.\"\"\"\n",
    "    return abs(x[0] - 0.) < np.finfo(float).eps  # type: ignore[no-any-return]\n",
    "\n",
    "\n",
    "def top(x: np.typing.NDArray[np.float64]) -> np.typing.NDArray[np.bool_]:\n",
    "    \"\"\"Condition that defines the top boundary.\"\"\"\n",
    "    return abs(x[1] - 1.) < np.finfo(float).eps  # type: ignore[no-any-return]\n",
    "\n",
    "\n",
    "def right(x: np.typing.NDArray[np.float64]) -> np.typing.NDArray[np.bool_]:\n",
    "    \"\"\"Condition that defines the right boundary.\"\"\"\n",
    "    return abs(x[0] - 1.) < np.finfo(float).eps  # type: ignore[no-any-return]\n",
    "\n",
    "\n",
    "boundaries_entities = dict()\n",
    "boundaries_values = dict()\n",
    "for (boundary, boundary_id) in zip((bottom, left, top, right), (1, 2, 3, 4)):\n",
    "    boundaries_entities[boundary_id] = dolfinx.mesh.locate_entities_boundary(\n",
    "        mesh, mesh.topology.dim - 1, boundary)\n",
    "    boundaries_values[boundary_id] = np.full(\n",
    "        boundaries_entities[boundary_id].shape, boundary_id, dtype=np.int32)\n",
    "boundaries_entities_unsorted = np.hstack(list(boundaries_entities.values()))\n",
    "boundaries_values_unsorted = np.hstack(list(boundaries_values.values()))\n",
    "boundaries_entities_argsort = np.argsort(boundaries_entities_unsorted)\n",
    "boundaries_entities_sorted = boundaries_entities_unsorted[boundaries_entities_argsort]\n",
    "boundaries_values_sorted = boundaries_values_unsorted[boundaries_entities_argsort]\n",
    "boundaries = dolfinx.mesh.meshtags(\n",
    "    mesh, mesh.topology.dim - 1,\n",
    "    boundaries_entities_sorted, boundaries_values_sorted)\n",
    "boundaries.name = \"boundaries\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7697d7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "boundaries_2 = boundaries.indices[boundaries.values == 2]\n",
    "boundaries_4 = boundaries.indices[boundaries.values == 4]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df464c5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define associated measures\n",
    "ds = ufl.Measure(\"ds\", subdomain_data=boundaries)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea07a303",
   "metadata": {},
   "outputs": [],
   "source": [
    "viskex.dolfinx.plot_mesh(mesh)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d22aff2",
   "metadata": {},
   "outputs": [],
   "source": [
    "viskex.dolfinx.plot_mesh_tags(mesh, boundaries, \"boundaries\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fee2bc92",
   "metadata": {},
   "source": [
    "### Function spaces"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "633b22e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "Y = dolfinx.fem.functionspace(mesh, (\"Lagrange\", 2))\n",
    "U = dolfinx.fem.functionspace(mesh, (\"Lagrange\", 2))\n",
    "Q = Y.clone()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c067fa2",
   "metadata": {},
   "source": [
    "### Restrictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae8218f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "dofs_Y = np.arange(0, Y.dofmap.index_map.size_local + Y.dofmap.index_map.num_ghosts)\n",
    "dofs_U = dolfinx.fem.locate_dofs_topological(U, boundaries.dim, boundaries_2)\n",
    "dofs_Q = dofs_Y\n",
    "restriction_Y = multiphenicsx.fem.DofMapRestriction(Y.dofmap, dofs_Y)\n",
    "restriction_U = multiphenicsx.fem.DofMapRestriction(U.dofmap, dofs_U)\n",
    "restriction_Q = multiphenicsx.fem.DofMapRestriction(Q.dofmap, dofs_Q)\n",
    "restriction = [restriction_Y, restriction_U, restriction_Q]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a3553389",
   "metadata": {},
   "source": [
    "### Trial and test functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "892f0cdc",
   "metadata": {},
   "outputs": [],
   "source": [
    "(y, u, p) = (ufl.TrialFunction(Y), ufl.TrialFunction(U), ufl.TrialFunction(Q))\n",
    "(z, v, q) = (ufl.TestFunction(Y), ufl.TestFunction(U), ufl.TestFunction(Q))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9a3917c",
   "metadata": {},
   "source": [
    " ### Problem data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e4f118f",
   "metadata": {},
   "outputs": [],
   "source": [
    "alpha = 1.e-5\n",
    "y_d = 1.\n",
    "epsilon = 1.e-1\n",
    "beta = ufl.as_vector((-1., -2.))\n",
    "sigma = 1.\n",
    "ff = 1.\n",
    "bc0 = petsc4py.PETSc.ScalarType(0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a6bf6d6",
   "metadata": {},
   "source": [
    "### Optimality conditions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d339c6a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "state_operator = (epsilon * ufl.inner(ufl.grad(y), ufl.grad(q)) * ufl.dx\n",
    "                  + ufl.inner(ufl.dot(beta, ufl.grad(y)), q) * ufl.dx + sigma * ufl.inner(y, q) * ufl.dx)\n",
    "adjoint_operator = (epsilon * ufl.inner(ufl.grad(p), ufl.grad(z)) * ufl.dx\n",
    "                    - ufl.inner(ufl.dot(beta, ufl.grad(p)), z) * ufl.dx + sigma * ufl.inner(p, z) * ufl.dx)\n",
    "a = [[ufl.inner(y, z) * ufl.dx, None, adjoint_operator],\n",
    "     [None, alpha * ufl.inner(u, v) * ds(2), - ufl.inner(p, v) * ds(2)],\n",
    "     [state_operator, - ufl.inner(u, q) * ds(2), None]]\n",
    "f = [ufl.inner(y_d, z) * ufl.dx,\n",
    "     None,\n",
    "     ufl.inner(ff, q) * ufl.dx]\n",
    "a[2][2] = dolfinx.fem.Constant(mesh, petsc4py.PETSc.ScalarType(0)) * ufl.inner(p, q) * ufl.dx\n",
    "f[1] = ufl.inner(dolfinx.fem.Constant(mesh, petsc4py.PETSc.ScalarType(0)), v) * ufl.dx\n",
    "a_cpp = dolfinx.fem.form(a)\n",
    "f_cpp = dolfinx.fem.form(f)\n",
    "bdofs_Y_4 = dolfinx.fem.locate_dofs_topological(Y, mesh.topology.dim - 1, boundaries_4)\n",
    "bdofs_Q_4 = dolfinx.fem.locate_dofs_topological(Q, mesh.topology.dim - 1, boundaries_4)\n",
    "bc = [dolfinx.fem.dirichletbc(bc0, bdofs_Y_4, Y),\n",
    "      dolfinx.fem.dirichletbc(bc0, bdofs_Q_4, Q)]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4143e08",
   "metadata": {},
   "source": [
    "### Solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57c5f349",
   "metadata": {},
   "outputs": [],
   "source": [
    "(y, u, p) = (dolfinx.fem.Function(Y), dolfinx.fem.Function(U), dolfinx.fem.Function(Q))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "915cd82f",
   "metadata": {},
   "source": [
    "### Cost functional"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53a9843a",
   "metadata": {},
   "outputs": [],
   "source": [
    "J = 0.5 * ufl.inner(y - y_d, y - y_d) * ufl.dx + 0.5 * alpha * ufl.inner(u, u) * ds(2)\n",
    "J_cpp = dolfinx.fem.form(J)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f4e0a44",
   "metadata": {},
   "source": [
    "### Uncontrolled functional value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a20f249c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Extract state forms from the optimality conditions\n",
    "a_state = ufl.replace(a[2][0], {q: z})\n",
    "f_state = ufl.replace(f[2], {q: z})\n",
    "a_state_cpp = dolfinx.fem.form(a_state)\n",
    "f_state_cpp = dolfinx.fem.form(f_state)\n",
    "bc_state = [bc[0]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37ea8961",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Assemble the linear system for the state\n",
    "A_state = dolfinx.fem.petsc.assemble_matrix(a_state_cpp, bcs=bc_state)\n",
    "A_state.assemble()\n",
    "F_state = dolfinx.fem.petsc.assemble_vector(f_state_cpp)\n",
    "dolfinx.fem.petsc.apply_lifting(F_state, [a_state_cpp], [bc_state])\n",
    "F_state.ghostUpdate(addv=petsc4py.PETSc.InsertMode.ADD, mode=petsc4py.PETSc.ScatterMode.REVERSE)\n",
    "dolfinx.fem.petsc.set_bc(F_state, bc_state)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b59977fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solve\n",
    "ksp = petsc4py.PETSc.KSP()\n",
    "ksp.create(mesh.comm)\n",
    "ksp.setOperators(A_state)\n",
    "ksp.setType(\"preonly\")\n",
    "ksp.getPC().setType(\"lu\")\n",
    "ksp.getPC().setFactorSolverType(\"mumps\")\n",
    "ksp.setFromOptions()\n",
    "ksp.solve(F_state, y.x.petsc_vec)\n",
    "y.x.petsc_vec.ghostUpdate(addv=petsc4py.PETSc.InsertMode.INSERT, mode=petsc4py.PETSc.ScatterMode.FORWARD)\n",
    "ksp.destroy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5005fc48",
   "metadata": {},
   "outputs": [],
   "source": [
    "J_uncontrolled = mesh.comm.allreduce(dolfinx.fem.assemble_scalar(J_cpp), op=mpi4py.MPI.SUM)\n",
    "print(\"Uncontrolled J =\", J_uncontrolled)\n",
    "assert np.isclose(J_uncontrolled, 0.23058804)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ba0f335",
   "metadata": {},
   "outputs": [],
   "source": [
    "viskex.dolfinx.plot_scalar_field(y, \"uncontrolled state\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7e09c60b",
   "metadata": {},
   "source": [
    "### Optimal control"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0d36f8c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Assemble the block linear system for the optimality conditions\n",
    "A = multiphenicsx.fem.petsc.assemble_matrix_block(a_cpp, bcs=bc, restriction=(restriction, restriction))\n",
    "A.assemble()\n",
    "F = multiphenicsx.fem.petsc.assemble_vector_block(f_cpp, a_cpp, bcs=bc, restriction=restriction)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2f8aa0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solve\n",
    "yup = multiphenicsx.fem.petsc.create_vector_block(f_cpp, restriction=restriction)\n",
    "ksp = petsc4py.PETSc.KSP()\n",
    "ksp.create(mesh.comm)\n",
    "ksp.setOperators(A)\n",
    "ksp.setType(\"preonly\")\n",
    "ksp.getPC().setType(\"lu\")\n",
    "ksp.getPC().setFactorSolverType(\"mumps\")\n",
    "ksp.setFromOptions()\n",
    "ksp.solve(F, yup)\n",
    "yup.ghostUpdate(addv=petsc4py.PETSc.InsertMode.INSERT, mode=petsc4py.PETSc.ScatterMode.FORWARD)\n",
    "ksp.destroy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e49be85",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Split the block solution in components\n",
    "with multiphenicsx.fem.petsc.BlockVecSubVectorWrapper(yup, [Y.dofmap, U.dofmap, Q.dofmap], restriction) as yup_wrapper:\n",
    "    for yup_wrapper_local, component in zip(yup_wrapper, (y, u, p)):\n",
    "        with component.x.petsc_vec.localForm() as component_local:\n",
    "            component_local[:] = yup_wrapper_local"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4bec2de",
   "metadata": {},
   "outputs": [],
   "source": [
    "J_controlled = mesh.comm.allreduce(dolfinx.fem.assemble_scalar(J_cpp), op=mpi4py.MPI.SUM)\n",
    "print(\"Optimal J =\", J_controlled)\n",
    "assert np.isclose(J_controlled, 0.21175842)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73a709f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "viskex.dolfinx.plot_scalar_field(y, \"state\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25b57a5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "viskex.dolfinx.plot_scalar_field(u, \"control\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37fd863a",
   "metadata": {},
   "outputs": [],
   "source": [
    "viskex.dolfinx.plot_scalar_field(p, \"adjoint\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython"
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
